import os

import cv2
import numpy as np
from PIL import Image
from typing import Tuple

from get_project_path import project_path
from utils.logger import GetLogger

logger = GetLogger.get_logger()
class ImageComparator:

    @staticmethod
    def compare_images(
        current_image_path: str,
        baseline_image_path: str,
        threshold: float = 0.95,
        diff_image_path: str = None,
    ) -> Tuple[bool, float, np.ndarray]:
        """
        对比两张图片的相似度，并生成差异图。

        :param current_image_path: 当前图片路径
        :param baseline_image_path: 基线图片路径
        :param threshold: 相似度阈值（0-1）
        :param diff_image_path: 差异图保存路径（可选）
        :return: (是否通过, 相似度, 差异图数组)
        """
        # 如果基线图片不存在，将当前图片保存为基线图片
        if not os.path.exists(baseline_image_path):
            logger.info(f"基线图片不存在，将当前图片保存为基线图片: {baseline_image_path}")
            os.makedirs(os.path.dirname(baseline_image_path), exist_ok=True)
            Image.open(current_image_path).save(baseline_image_path)
            return True, 1.0, np.zeros((1, 1))  # 首次运行直接返回通过

        # 读取图片
        current_image = cv2.imread(current_image_path, cv2.IMREAD_COLOR)
        baseline_image = cv2.imread(baseline_image_path, cv2.IMREAD_COLOR)

        # 检查图片是否读取成功
        if current_image is None:
            raise ValueError(f"无法读取当前图片: {current_image_path}")
        if baseline_image is None:
            raise ValueError(f"无法读取基线图片: {baseline_image_path}")

        # 确保图片尺寸一致
        if current_image.shape != baseline_image.shape:
            raise ValueError("图片尺寸不一致，无法比较")

        # 转换为灰度图
        current_gray = cv2.cvtColor(current_image, cv2.COLOR_BGR2GRAY)
        baseline_gray = cv2.cvtColor(baseline_image, cv2.COLOR_BGR2GRAY)

        # 计算结构相似性 (SSIM)
        ssim_score, diff = ImageComparator._calculate_ssim(current_gray, baseline_gray)

        # 将差异图归一化到 0-255 范围
        diff_normalized = (diff * 255).astype("uint8")

        # 如果需要保存差异图
        if diff_image_path:
            diff_image = Image.fromarray(diff_normalized)
            diff_image.save(diff_image_path)

        # 返回结果
        is_pass = ssim_score >= threshold
        logger.info(f"SSIM 分数: {ssim_score}")
        return is_pass, ssim_score, diff_normalized

    @staticmethod
    def _calculate_ssim(image1: np.ndarray, image2: np.ndarray) -> Tuple[float, np.ndarray]:
        """
        计算两张图片的结构相似性 (SSIM)。

        :param image1: 图片1
        :param image2: 图片2
        :return: (SSIM 分数, 差异图)
        """
        # SSIM 参数
        C1 = (0.01 * 255) ** 2
        C2 = (0.03 * 255) ** 2

        # 计算均值、方差和协方差
        mu1 = cv2.GaussianBlur(image1, (11, 11), 1.5)
        mu2 = cv2.GaussianBlur(image2, (11, 11), 1.5)
        mu1_sq = mu1 ** 2
        mu2_sq = mu2 ** 2
        mu1_mu2 = mu1 * mu2

        sigma1_sq = cv2.GaussianBlur(image1 ** 2, (11, 11), 1.5) - mu1_sq
        sigma2_sq = cv2.GaussianBlur(image2 ** 2, (11, 11), 1.5) - mu2_sq
        sigma12 = cv2.GaussianBlur(image1 * image2, (11, 11), 1.5) - mu1_mu2

        # 计算 SSIM
        ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / (
            (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
        )
        ssim_score = ssim_map.mean()

        return ssim_score, ssim_map

    @staticmethod
    def compare(
        current_image_path: str,
        baseline_image_path: str,
        threshold: float = 0.95,
        diff_image_path: str = None,
    ) -> bool:
        """
        简化版的图片对比方法，仅返回是否通过。

        :param current_image_path: 当前图片路径
        :param baseline_image_path: 基线图片路径
        :param threshold: 相似度阈值（0-1）
        :param diff_image_path: 差异图保存路径（可选）
        :return: 是否通过
        """
        is_pass, _, _ = ImageComparator.compare_images(
            current_image_path, baseline_image_path, threshold, diff_image_path
        )
        return is_pass
if __name__ == '__main__':
    print(ImageComparator.compare(os.path.join(project_path, "images", "GR-v8_2_1-手机填写页.png"),
                                  os.path.join(project_path, "images", "GR2024-v8_2_1-手机填写页.png")
                                  ))